import json
from io import BytesIO

from rest_framework.decorators import action

from backend.utils.base_viewsets import ExportableViewSet
from backend.utils.menu_tree import DepartmentTree
from backend.utils.response import ApiResponse
from department.models import Department
from department.serializers import DepartmentCreateSerializer, DepartmentExportSerializer


class DepartmentModelViewSet(ExportableViewSet):
    serializer_class = {
        'default': DepartmentCreateSerializer,
        'export': DepartmentExportSerializer,
    }
    excel_headers = ['部门名称', '地址', '状态', '备注', '创建时间', '更新时间', '上级部门', ]
    file_name = '部门列表.xlsx'

    def get_queryset(self):
        department_name = self.request.query_params.get('name')
        status = self.request.query_params.get('status')
        queryset = Department.advanced_objects.all()
        if department_name and self.action == 'list':
            queryset = queryset.filter(name__icontains=department_name)
        if status and self.action == 'list':
            status = json.loads(status)
            queryset = queryset.filter(status=status)
        return queryset.order_by('-update_time', 'id')

    def list(self, request, *args, **kwargs):
        page = int(self.request.query_params.get('page', 1))
        page_size = int(self.request.query_params.get('page_size', -1))
        response = super(DepartmentModelViewSet, self).list(request, *args, **kwargs)
        data = DepartmentTree.make_department_tree(response.data.get('data'), None)
        total = len(data)
        start = (page - 1) * page_size
        end = page * page_size
        if page_size == -1:
            end = total
        data = data[start:end]
        return ApiResponse(data=data, total=total)

    def destroy(self, request, *args, **kwargs):
        department = self.get_object()
        if department.children.count() > 0:
            return ApiResponse(code=400, message='该部门下有子部门无法删除')
        # TODO 关联用户表后,查询部门下是否有用户,如果有用户,则不能删除
        return super(DepartmentModelViewSet, self).destroy(request, *args, **kwargs)
